import dotenv
import requests
from langchain_community.tools import GoogleSerperRun
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_core.messages import ToolMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.tools import BaseTool
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

import json
import os
from typing import Type, Any
dotenv.load_dotenv()

class GaodeWeatherSchema(BaseModel):
    city:str = Field(description="需要查询天气的城市，如：广州")

class GoogleSerperSchema(BaseModel):
    query:str = Field(description="执行谷歌搜索的查询语句")
class GaodeWeatherTool(BaseTool):
    """根据传入的城市名称运行调用API 获取传入城市的天气预报信息"""
    name:str = "gaode_weather_tool"
    description:str = "当你想查询天气都时候可以用这个工具"
    args_schema: Type[BaseModel] = GaodeWeatherSchema

    def _run(self, *args:Any, **kwargs:Any) -> Any:
        try:
            gaode_api_key = os.getenv("GAODE_API_KEY")
            gaode_api_url = os.getenv("GAODE_BASE_URL")
            if not gaode_api_key or not gaode_api_url:
                return f"请配置正确的高德API KEY 和 URL"
            else:
                #1、从参数中获取城市
                city = kwargs.get("city")
                #2、行政区域查询
                session = requests.session()
                city_response = session.request(
                    method = "GET",
                    url = f"{gaode_api_url}/config/district?key={gaode_api_key}&keywords={city}&subdistrict=0",
                    headers = {"Content-Type":"application/json;charset=utf-8"}
                )
                city_response.raise_for_status()
                city_data =city_response.json()
                print(city_data)
                if city_data.get("info") == "OK":
                    ad_code = city_data.get("districts")[0]["adcode"]
                    weather_info = session.request(
                        method="GET",
                        url=f"{gaode_api_url}/weather/weatherInfo?key={gaode_api_key}&city={ad_code}&extensions=all",
                    )
                    print(weather_info)
                    weather_info.raise_for_status()
                    weather_data = weather_info.json()
                    print(weather_data)
                    if weather_data.get("info") == "OK":
                        return json.dumps(weather_data)
                return f"获取{city}天气预报失败"
        except Exception as e:
            return f"获取天气失败"

# 定义工具列表
gaode_weather = GaodeWeatherTool()
google_serper = GoogleSerperRun(
    name = "google_serper",
    description = (
        "一个低成本的谷歌搜索工具"
    ),
    api_wrapper = GoogleSerperAPIWrapper()
)

tool_dict = {
    gaode_weather.name:gaode_weather,
    google_serper.name:google_serper
}

tools = [tool for tool in tool_dict.values()]

# 创建prompt
prompt = ChatPromptTemplate.from_messages(
    [
        ("system","你是由OpenAI研发的机器人"),
        ("human","{query}")
    ]
)

# 创建大模型并绑定工具
llm = ChatOpenAI(model="gpt-4o-mini")
llm_with_tool = llm.bind_tools(tools=tools)#绑定两个工具到大模型中，一个是serper 一个是 weather

# 创建应用
chain = {"query":RunnablePassthrough()} | prompt | llm_with_tool
# 解析输出
query = ("请问广州天气如何")
resp = chain.invoke(query)#将用户的提问提交大模型
tool_calls = resp.tool_calls #将返回内容中的工具信息提取出来

#判断工具是否有正常输出结果
if len(tool_calls) <=0:  #==0  说明没有工具调用信息
    print("生成的内容是",resp.content)
else:  #>0  说明有工具调用信息
    # 将历史消息、人类消息、AI消息结合
    messages = prompt.invoke(query).to_messages()#组合整个聊天过程
    messages.append(resp)

#     循环遍历所有工具调用
    for tool_call in tool_calls:#解析工具调用信息
        tool = tool_dict.get(tool_call.get('name'))#根据大模型返回工具名称来查找工具的实例
        print('正在执行这个工具：',tool.name)
        id = tool_call.get('id')#获取工具调用的id 为后续构造工具信息提供参数
        content = tool.invoke(tool_call.get("args"))#从大模型返回的内容解析工具调用参数并通过工具实例工具调用
        print("工具输出",content)
        messages.append(ToolMessage(#工具输出后把工具调用的结果打包成工具消息
            content = content,
            tool_call_id = id
        ))
    print('输出内容：',llm.invoke(messages))#大模型整理工具消息并返回